#include "checksum_util.h"
#include "common.h"
#include <arpa/inet.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>

// 计算校验和
unsigned short calc_checksum(unsigned short* buffer, int size)
{
    unsigned long cksum = 0;
    while(size>1)
    {
        cksum += *buffer++;
        size -= sizeof(unsigned short);
    }
    if(size)
    {
        cksum += *(unsigned char*)buffer;
    }
    cksum = (cksum>>16) + (cksum&0xffff);
    cksum += (cksum>>16);
    return (unsigned short)(~cksum);
}

// 计算TCP校验和
void calc_tcp_checksum(char* raw_buffer){
    if (raw_buffer == NULL){
        return;
    }

    // RAW包头
    struct iphdr * ih_ptr = (struct iphdr *)raw_buffer;

    // 头长度
    __u16 ihl = (unsigned char)(ih_ptr->ihl) * 4;
    // IP包总长度
    __u16 tot_len = ntohs(ih_ptr->tot_len);
    // 传输层长度
    __u16 t_layer_len = tot_len - ihl;
    // 分配长度
    __u16 alloc_size = t_layer_len + sizeof(psdhdr);
    // 传输层头
    struct tcphdr* t_layer_header = (struct tcphdr*)(raw_buffer + ihl);

    // 伪首部指针
    psdhdr* ph_ptr = (psdhdr*)malloc(alloc_size);
    if(!ph_ptr) return;

    memset(ph_ptr,0,alloc_size);

    //填充伪首部
    ph_ptr->dest=ih_ptr->daddr;
    ph_ptr->source=ih_ptr->saddr;
    ph_ptr->mbz=0;
    ph_ptr->protocol=0x06;
    ph_ptr->len=htons(t_layer_len);

    //计算TCP校验和
    t_layer_header->check = 0;
    memcpy((unsigned char *)ph_ptr+sizeof(psdhdr), (unsigned char *)t_layer_header, t_layer_len);
    unsigned short new_check = calc_checksum((unsigned short *)ph_ptr, alloc_size);
    t_layer_header->check = new_check;
    //LOGI("tcp new_check=%u",new_check);

    // 计算IP校验和
    ih_ptr->check = 0;
    new_check = calc_checksum((unsigned short*)raw_buffer, ihl);
    ih_ptr->check = new_check;
    //LOGI("ip new_check=%u",new_check);

    free(ph_ptr);
}

void calc_udp_checksum(char* raw_buffer){
    if (raw_buffer == NULL){
        return;
    }

    // RAW包头
    struct iphdr * ih_ptr = (struct iphdr *)raw_buffer;

    // 头长度
    __u16 ihl = (unsigned char)(ih_ptr->ihl) * 4;
    // IP包总长度
    __u16 tot_len = ntohs(ih_ptr->tot_len);
    // 传输层长度
    __u16 t_layer_len = tot_len - ihl;
    // 分配长度
    __u16 alloc_size = t_layer_len + sizeof(psdhdr);
    // 传输层头
    struct udphdr* t_layer_header = (struct udphdr*)(raw_buffer + ihl);

    // 伪首部指针
    psdhdr* ph_ptr = (psdhdr*)malloc(alloc_size);
    if(!ph_ptr) return;

    memset(ph_ptr,0,alloc_size);

    //填充伪首部
    ph_ptr->dest=ih_ptr->daddr;
    ph_ptr->source=ih_ptr->saddr;
    ph_ptr->mbz=0;
    ph_ptr->protocol=0x11;
    ph_ptr->len=htons(t_layer_len);

    //计算UDP校验和
    t_layer_header->check = 0;
    memcpy((unsigned char *)ph_ptr+sizeof(psdhdr), (unsigned char *)t_layer_header, t_layer_len);
    unsigned short new_check = calc_checksum((unsigned short *)ph_ptr, alloc_size);
    t_layer_header->check = new_check;
    //LOGI("udp new_check=%u",new_check);

    // 计算IP校验和
    ih_ptr->check = 0;
    new_check = calc_checksum((unsigned short*)raw_buffer, ihl);
    ih_ptr->check = new_check;
    //LOGI("ip new_check=%u",new_check);

    free(ph_ptr);
}

bool check_tcp_checksum(char* raw_buffer){
    if (raw_buffer == NULL){
        return false;
    }

    // RAW包头
    struct iphdr * ih_ptr = (struct iphdr *)raw_buffer;

    // 头长度
    __u16 ihl = (unsigned char)(ih_ptr->ihl) * 4;
    // IP包总长度
    __u16 tot_len = ntohs(ih_ptr->tot_len);
    // 传输层长度
    __u16 t_layer_len = tot_len - ihl;
    // 分配长度
    __u16 alloc_size = t_layer_len + sizeof(psdhdr);
    // 传输层头
    struct tcphdr* t_layer_header = (struct tcphdr*)(raw_buffer + ihl);

    // 伪首部指针
    psdhdr* ph_ptr = (psdhdr*)malloc(alloc_size);
    if(!ph_ptr) return false;

    memset(ph_ptr,0,alloc_size);

    //填充伪首部
    ph_ptr->dest=ih_ptr->daddr;
    ph_ptr->source=ih_ptr->saddr;
    ph_ptr->mbz=0;
    ph_ptr->protocol=0x06;
    ph_ptr->len=htons(t_layer_len);

    //计算TCP校验和
    memcpy((unsigned char *)ph_ptr+sizeof(psdhdr), (unsigned char *)t_layer_header, t_layer_len);
    unsigned short new_check_tcp = calc_checksum((unsigned short *)ph_ptr, alloc_size);
    //LOGI("check tcp checksum=%u",new_check_tcp);

    // 计算IP校验和
    unsigned short new_check_ip = calc_checksum((unsigned short*)raw_buffer, ihl);
    //LOGI("check ip checksum=%u", new_check_ip);

    free(ph_ptr);

    return (!new_check_tcp && !new_check_ip);
}

bool check_udp_checksum(char* raw_buffer){
    if (raw_buffer == NULL){
        return false;
    }

    // RAW包头
    struct iphdr * ih_ptr = (struct iphdr *)raw_buffer;

    // 头长度
    __u16 ihl = (unsigned char)(ih_ptr->ihl) * 4;
    // IP包总长度
    __u16 tot_len = ntohs(ih_ptr->tot_len);
    // 传输层长度
    __u16 t_layer_len = tot_len - ihl;
    // 分配长度
    __u16 alloc_size = t_layer_len + sizeof(psdhdr);
    // 传输层头
    struct udphdr* t_layer_header = (struct udphdr*)(raw_buffer + ihl);

    // 伪首部指针
    psdhdr* ph_ptr = (psdhdr*)malloc(alloc_size);
    if(!ph_ptr) return false;

    memset(ph_ptr,0,alloc_size);

    //填充伪首部
    ph_ptr->dest=ih_ptr->daddr;
    ph_ptr->source=ih_ptr->saddr;
    ph_ptr->mbz=0;
    ph_ptr->protocol=0x11;
    ph_ptr->len=htons(t_layer_len);

    //计算UDP校验和
    memcpy((unsigned char *)ph_ptr+sizeof(psdhdr), (unsigned char *)t_layer_header, t_layer_len);
    unsigned short new_check_udp = calc_checksum((unsigned short *)ph_ptr, alloc_size);
    //LOGI("check udp checksum=%u",new_check_udp);

    // 计算IP校验和
    unsigned short new_check_ip = calc_checksum((unsigned short*)raw_buffer, ihl);
    //LOGI("check ip checksum=%u",new_check_ip);

    free(ph_ptr);

    return (!new_check_udp && !new_check_ip);
}
